import sqlite3
import csv
import os

conn = sqlite3.connect("E:\\immunedb\\immuno.db")
cursor = conn.cursor()
#hla名称，下面的这个只是示例，跑的时候换其他的就可以
try:
    specified_hlas = [
        "HLA-B*53:01"
    ]
    print(f"共处理 {len(specified_hlas)} 种HLA")

    # 步骤1：创建临时表（新增exp_level_num列存储映射后的数值）
    cursor.execute("""
        CREATE TEMP TABLE IF NOT EXISTS temp_joined_data (
            hla TEXT,
            tissue TEXT,
            probability REAL,
            exp_level_num INTEGER  -- 存储exp_level映射数值：3=High/2=Medium/1=Low/0=其他
        );
    """)

    # 步骤2：分批插入数据（含exp_level数值映射，保留所有表达水平）
    print("\n=== 步骤2：分批插入关联数据（含exp_level数值映射） ===")
    for i, hla in enumerate(specified_hlas, 1):
        cursor.execute(f"""
            SELECT COUNT(*) 
            FROM predictions p
            JOIN peptides pe ON p.peptide_id = pe.peptide_id
            JOIN protein_gene pg ON pe.protein_id = pg.protein_id
            JOIN protein_exp pe_exp ON pg.geneName = pe_exp.geneName
            WHERE p.hla = '{hla}';
        """)
        count = cursor.fetchone()[0]
        estimated_size_mb = count * 0.12 / 1024
        print(f"- 第{i}/{len(specified_hlas)}种：{hla}，共 {count} 条数据，预估 {estimated_size_mb:.2f} MB")

        if count > 0:
            cursor.execute(f"""
                INSERT INTO temp_joined_data (hla, tissue, probability, exp_level_num)
                SELECT 
                    p.hla, 
                    pe_exp.tissue, 
                    p.probability,
                    CASE 
                        WHEN pe_exp.exp_level = 'High' THEN 3
                        WHEN pe_exp.exp_level = 'Medium' THEN 2
                        WHEN pe_exp.exp_level = 'Low' THEN 1
                        ELSE 0
                    END AS exp_level_num
                FROM predictions p
                JOIN peptides pe ON p.peptide_id = pe.peptide_id
                JOIN protein_gene pg ON pe.protein_id = pg.protein_id
                JOIN protein_exp pe_exp ON pg.geneName = pe_exp.geneName
                WHERE p.hla = '{hla}';
            """)
            conn.commit()
            print(f"  已插入临时表（含exp_level数值映射）\n")
        else:
            print(f"  无数据，跳过\n")

    # 步骤3：按HLA+组织拆分统计（仅probability>0.5的肽段平均表达水平）
    print("\n=== 步骤3：按HLA+组织拆分统计（仅prob>0.5的肽段平均表达水平） ===")
    # 核心修改：CSV文件名改为 protein_hla_tissue.csv
    output_path = "E:/immunedb/protein_hla_tissue.csv"
    with open(output_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "HLA", "组织类型", "概率>0.5的数量", "概率<0.5的数量",
            "平均表达水平（仅prob>0.5，3=High/2=Medium/1=Low）"
        ])

        for hla in specified_hlas:
            cursor.execute(f"""
                SELECT DISTINCT tissue 
                FROM temp_joined_data 
                WHERE hla = '{hla}';
            """)
            tissues = [t[0] for t in cursor.fetchall()]

            for tissue in tissues:
                cursor.execute(f"""
                    SELECT 
                        '{hla}' AS hla,
                        '{tissue}' AS tissue,
                        SUM(CASE WHEN probability > 0.5 THEN 1 ELSE 0 END) AS count_gt_0_5,
                        SUM(CASE WHEN probability < 0.5 THEN 1 ELSE 0 END) AS count_lt_0_5,
                        ROUND(
                            CAST(
                                SUM(CASE WHEN probability > 0.5 THEN exp_level_num ELSE 0 END) 
                            AS REAL) 
                            / MAX(SUM(CASE WHEN probability > 0.5 THEN 1 ELSE 0 END), 1),
                            2
                        ) AS avg_exp_level_gt_0_5
                    FROM temp_joined_data
                    WHERE hla = '{hla}' AND tissue = '{tissue}';
                """)
                result = cursor.fetchone()
                if result:
                    writer.writerow(result)
                print(f"已统计：{hla} + {tissue} → 平均表达水平：{result[4]}")

    print(f"\n所有结果已保存到：{output_path}")
    print(f"输出说明：")
    print(f"  1. 平均表达水平：仅基于probability>0.5的肽段计算（exp_level数值总和 ÷ 该部分肽段数量）")
    print(f"  2. 数值映射规则：High→3、Medium→2、Low→1、其他→0")

except sqlite3.Error as e:
    print(f"\n数据库错误: {e}")
finally:
    cursor.close()
    conn.close()
    print("\n程序结束")